#!/usr/bin/env python3
"""
kmod MCP Server

面向 kmod 软件（内核模块管理工具）的 MCP 服务器，提供内核模块管理的命令行操作封装。

MCP tools 列表：
- kmod_list: 列出已加载的内核模块
- kmod_info: 显示指定内核模块的详细信息
- kmod_load: 加载指定的内核模块
- kmod_unload: 卸载指定的内核模块  
- kmod_dependencies: 显示模块的依赖关系
"""

from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json
from typing import List, Optional

mcp = FastMCP("kmod MCP Server")

@mcp.tool()
def kmod_list(ctx: Context) -> dict:
    """
    列出系统中所有已加载的内核模块
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令执行是否成功
        "command": str,       # 实际执行的命令
        "exit_code": int,     # 命令退出码
        "stdout": str,        # 命令标准输出
        "stderr": str         # 命令标准错误
    }
    """
    command = ["lsmod"]
    result = subprocess.run(command, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(command),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def kmod_info(ctx: Context, module_name: str) -> dict:
    """
    显示指定内核模块的详细信息
    
    参数：
    - module_name: 内核模块名称
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令执行是否成功
        "command": str,       # 实际执行的命令
        "exit_code": int,     # 命令退出码
        "stdout": str,        # 命令标准输出
        "stderr": str         # 命令标准错误
    }
    """
    command = ["modinfo", module_name]
    result = subprocess.run(command, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(command),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def kmod_load(ctx: Context, module_name: str) -> dict:
    """
    加载指定的内核模块
    
    参数：
    - module_name: 要加载的内核模块名称
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令执行是否成功
        "command": str,       # 实际执行的命令
        "exit_code": int,     # 命令退出码
        "stdout": str,        # 命令标准输出
        "stderr": str         # 命令标准错误
    }
    """
    command = ["modprobe", module_name]
    result = subprocess.run(command, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(command),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def kmod_unload(ctx: Context, module_name: str) -> dict:
    """
    卸载指定的内核模块
    
    参数：
    - module_name: 要卸载的内核模块名称
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令执行是否成功
        "command": str,       # 实际执行的命令
        "exit_code": int,     # 命令退出码
        "stdout": str,        # 命令标准输出
        "stderr": str         # 命令标准错误
    }
    """
    command = ["modprobe", "-r", module_name]
    result = subprocess.run(command, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(command),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

@mcp.tool()
def kmod_dependencies(ctx: Context, module_name: str) -> dict:
    """
    显示指定内核模块的依赖关系
    
    参数：
    - module_name: 内核模块名称
    
    返回 JSON 结构：
    {
        "success": bool,      # 命令执行是否成功
        "command": str,       # 实际执行的命令
        "exit_code": int,     # 命令退出码
        "stdout": str,        # 命令标准输出
        "stderr": str         # 命令标准错误
    }
    """
    command = ["modinfo", "-F", "depends", module_name]
    result = subprocess.run(command, capture_output=True, text=True)
    
    return {
        "success": result.returncode == 0,
        "command": " ".join(command),
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr
    }

if __name__ == "__main__":
    mcp.run()